<!DOCTYPE html>
<html lang="zh">
<head>
  <meta charset="UTF-8" />
  <title>MNIST 手写识别</title>
  <style>
    body {
      font-family: sans-serif;
      text-align: center;
    }
    canvas {
      border: 1px solid #000;
      background-color: #fff;
      cursor: crosshair;
    }
    #result {
      margin-top: 20px;
      font-size: 2em;
      color: green;
    }
    button {
      margin-top: 10px;
      padding: 10px 20px;
      font-size: 1em;
    }
  </style>
</head>
<body>
  <h1>手写数字识别</h1>
  <canvas id="drawingCanvas" width="280" height="280"></canvas>
  <br /><br />
  <button id="predictBtn">识别数字</button>
  <button id="clearBtn">清除画布</button>
  <div id="result">识别结果：-</div>

  <!-- 引入 TensorFlow.js -->
  <script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@4.10.0/dist/tf.min.js"></script>

  <script>
    const canvas = document.getElementById('drawingCanvas');
    const ctx = canvas.getContext('2d');
    let drawing = false;

    // 清除画布
    function clearCanvas() {
      ctx.fillStyle = 'white';
      ctx.fillRect(0, 0, canvas.width, canvas.height);
    }

    // 鼠标事件监听
    canvas.addEventListener('mousedown', () => (drawing = true));
    canvas.addEventListener('mouseup', () => (drawing = false));
    canvas.addEventListener('mouseout', () => (drawing = false));
    canvas.addEventListener('mousemove', draw);

    function draw(e) {
      if (!drawing) return;
      ctx.lineWidth = 12;
      ctx.lineCap = 'round';
      ctx.strokeStyle = 'black';

      const rect = canvas.getBoundingClientRect();
      const x = e.clientX - rect.left;
      const y = e.clientY - rect.top;

      ctx.lineTo(x, y);
      ctx.stroke();
      ctx.beginPath();
      ctx.moveTo(x, y);
    }

    // 边界框裁剪函数
    function getBoundingBox(imageData) {
      let minX = Infinity,
        minY = Infinity,
        maxX = -Infinity,
        maxY = -Infinity;
      const data = imageData.data;

      for (let y = 0; y < imageData.height; y++) {
        for (let x = 0; x < imageData.width; x++) {
          const idx = (y * imageData.width + x) * 4;
          if (data[idx] > 0x10) {
            minX = Math.min(minX, x);
            minY = Math.min(minY, y);
            maxX = Math.max(maxX, x);
            maxY = Math.max(maxY, y);
          }
        }
      }

      if (minX === Infinity) return null;
      return { minX, minY, maxX, maxY };
    }

    // 图像预处理函数
    function preprocessCanvasForModel(ctx) {
      const imgData = ctx.getImageData(0, 0, canvas.width, canvas.height);
      const imageData = new Uint8ClampedArray(imgData.data);

      // 二值化处理
      for (let i = 0; i < imageData.length; i += 4) {
        const avg = (imageData[i] + imageData[i + 1] + imageData[i + 2]) / 3;
        const threshold = avg > 0x10 ? 255 : 0;
        imageData[i] = imageData[i + 1] = imageData[i + 2] = threshold;
        imageData[i + 3] = 255;
      }

      const tempCanvas = document.createElement('canvas');
      const tempCtx = tempCanvas.getContext('2d');
      tempCanvas.width = canvas.width;
      tempCanvas.height = canvas.height;
      tempCtx.putImageData(new ImageData(imageData, canvas.width, canvas.height), 0, 0);

      // 获取边界框并裁剪
      const boundingBox = getBoundingBox(tempCtx.getImageData(0, 0, tempCanvas.width, tempCanvas.height));
      if (!boundingBox) return tf.zeros([1, 28, 28, 1]);

      const { minX, minY, maxX, maxY } = boundingBox;
      const cropWidth = maxX - minX;
      const cropHeight = maxY - minY;

      const croppedCanvas = document.createElement('canvas');
      croppedCanvas.width = cropWidth;
      croppedCanvas.height = cropHeight;
      const croppedCtx = croppedCanvas.getContext('2d');
      croppedCtx.drawImage(tempCanvas, minX, minY, cropWidth, cropHeight, 0, 0, cropWidth, cropHeight);

      // 缩放到 28x28
      const finalCanvas = document.createElement('canvas');
      finalCanvas.width = 28;
      finalCanvas.height = 28;
      const finalCtx = finalCanvas.getContext('2d');
      finalCtx.fillStyle = 'white';
      finalCtx.fillRect(0, 0, 28, 28);
      finalCtx.drawImage(croppedCanvas, 0, 0, 28, 28);

      // 转换为 tensor
      let tensor = tf.browser.fromPixels(finalCanvas, 1).toFloat();
      return tensor.div(tf.scalar(255)).reshape([1, 28, 28, 1]);
    }

    // 加载模型
    let model;
    (async function () {
      try {
        model = await tf.loadLayersModel('/models/mnist/my-mnist-model.json');
        console.log('模型加载完成');
      } catch (err) {
        console.error('模型加载失败:', err);
        alert('模型加载失败，请检查路径是否正确或是否已启动本地服务器');
      }
    })();

    // 识别按钮点击事件
    document.getElementById('predictBtn').addEventListener('click', async () => {
      if (!model) {
        alert('模型尚未加载，请稍等...');
        return;
      }

      const input = tf.tidy(() => {
        return preprocessCanvasForModel(ctx);
      });

      const prediction = model.predict(input);
      const predictedIndex = prediction.argMax(1).dataSync()[0];

      document.getElementById('result').innerText = `识别结果：${predictedIndex}`;
      input.dispose();
      prediction.dispose();
    });

    // 清除按钮点击事件
    document.getElementById('clearBtn').addEventListener('click', () => {
      clearCanvas();
      document.getElementById('result').innerText = '识别结果：-';
    });
  </script>
</body>
</html>